#!/usr/bin/env python
# encoding: utf-8
"""
@author: youfeng
@email: youfeng243@163.com
@license: Apache Licence
@file: import_company_data.py
@time: 2018/1/22 18:46
"""
import os
import sys

import click

sys.path.append('..')
sys.path.append('../..')
from common import util
from common.mongo import MongDb
from logger import Logger

MONGO_DB_SOURCE = {
    "host": "172.16.215.2",
    "port": 40042,
    "db": "company_data",
    "username": "work",
    "password": "haizhi",
}

# 日志模块
log = Logger('import_company_data.log').get_logger()

company_data_db = MongDb(MONGO_DB_SOURCE['host'], MONGO_DB_SOURCE['port'], MONGO_DB_SOURCE['db'],
                         MONGO_DB_SOURCE['username'],
                         MONGO_DB_SOURCE['password'], log=log)


@click.command()
@click.option('--table_name',
              default='beijing_shanghai_hubo_20180122',
              help='需要导入的表名')
@click.option('--file_path',
              default='/Users/you/project/new_gsxt_crawl/script/gsxt/shanghaibeijing.txt',
              help='需要导入的表名')
def main(table_name, file_path):
    result_list = []
    count = 0

    if not isinstance(table_name, basestring) or table_name == '':
        log.error("表名称不符合规范: table_name = {}".format(table_name))
        return

    if not os.path.exists(path=file_path):
        log.error("当前文件不存在: path = {}".format(file_path))
        return

    with open(file_path) as p_file:
        for line in p_file:
            count += 1
            company = line.strip().strip('\n').strip('\r')
            data = {
                '_id': company,
                '_in_time': util.get_now_time(),
            }
            result_list.append(data)
            if len(result_list) >= 500:
                company_data_db.insert_batch_data(table_name, result_list)
                del result_list[:]

            if count % 1000 == 0:
                log.info("当前导入进度: total = {}".format(count))

    if len(result_list) > 0:
        company_data_db.insert_batch_data(table_name, result_list)
        del result_list[:]

    log.info("数据导入完成...")


if __name__ == '__main__':
    main()
